#!/bin/python3

import os
import json
import glob
import cv2 as cv

dataset_root = "/home/kx/Documents/datasets/retail_product/data"
test_image_path = os.path.join(dataset_root, "test2019")
train_image_path = os.path.join(dataset_root, "train2019")
val_image_path = os.path.join(dataset_root, "val2019")

train_anno_json = os.path.join(dataset_root, "instances_train2019.json")
test_anno_json = os.path.join(dataset_root, "instances_test2019.json")
val_anno_json = os.path.join(dataset_root, "instances_val2019.json")

output_dataset_root = os.path.join(dataset_root, "yolo")
output_train_path = os.path.join(output_dataset_root, "train")
output_test_path = os.path.join(output_dataset_root, "test")
output_val_path = os.path.join(output_dataset_root, "val")

def ParseAnno(json_file: str):
    anno_json = None
    with open(json_file, 'r') as fp:
        anno_json = json.load(fp)
    images = anno_json["images"]
    annotations = anno_json['annotations']
    return images, annotations

def doubleStr2Int(data: str):
    return int(float(data))

if __name__ == '__main__':
    if not os.path.exists(output_dataset_root):
        os.makedirs(output_dataset_root)
    if not os.path.exists(output_train_path):
        os.makedirs(output_train_path)
    if not os.path.exists(output_val_path):
        os.makedirs(output_val_path)
    if not os.path.exists(output_test_path):
        os.makedirs(output_test_path)

    # images, annotations = ParseAnno(test_anno_json)
    images, annotations = ParseAnno(train_anno_json)
    # print("images: ", images)
    images_direct = {}
    for image in images:
        file_name = image['file_name']
        id = image['id']
        images_direct[id] = file_name

    # print("annotations: ", annotations)
    # print("images_direct: ", images_direct)
    classes_count = {}
    file_annos = {}
    for anno in annotations:
        bbox = anno['bbox']
        category_id = anno['category_id']
        if category_id not in classes_count.keys():
            classes_count[category_id] = 1
        else:
            classes_count[category_id] += 1
        image_id = anno['image_id']
        file_name = images_direct[image_id]
        if file_name in file_annos.keys():
            file_annos[file_name].append({'bbox': bbox, 'id': category_id})
        else:
            file_annos[file_name] = []
            file_annos[file_name].append({'bbox': bbox, 'id': category_id})
    
    print("classes_count: ", classes_count)

    import matplotlib.pyplot as plt
    plt.bar(range(len(classes_count)), list(classes_count.values()), align='center')
    plt.show()


    ''' trans 2 txt
    for anno in file_annos:
        anno_txt_name = anno.split(".")[0] + ".txt"
        # anno_txt_name = os.path.join(output_test_path, anno_txt_name)
        anno_txt_name = os.path.join(output_val_path, anno_txt_name)
        # print("anno_txt_name: ", anno_txt_name)
        with open(anno_txt_name, "w")  as txt:
            lines = []
            for record in file_annos[anno]:
                # print("record: ", record)
                line  = str(record['id']) + ' ' + str(record['bbox'][0]) + ' ' + str(record['bbox'][1]) + ' ' + str(record['bbox'][2]) + ' ' + str(record['bbox'][3]) + '\n'
                # print("line: ", line)
                lines.append(line)
            txt.writelines(lines)
    '''
    # print("output_test_path: ", output_test_path + "/.txt")
    # for txt_file in glob.glob(output_test_path + "/*.txt"):
    #     # print("txt_file: ", txt_file)
    #     image_file = txt_file.split("/")[-1].split(".")[0] + ".jpg"
    #     image_path = os.path.join(test_image_path, image_file)
    #     # print("image_path: ", image_path)
    #     with open(txt_file, 'r') as an:
    #         lines = an.readlines()
    #         # print("lines: ", lines)
    #         image = cv.imread(image_path)
    #         for line in lines:
    #             line = line.strip().split(' ')
    #             # print("line: ", line)
    #             ltx = doubleStr2Int(line[1])
    #             lty = doubleStr2Int(line[2])
    #             rdx = ltx + doubleStr2Int(line[3])
    #             rdy = lty + doubleStr2Int(line[4])
    #             id = doubleStr2Int(line[0])
    #             cv.putText(image, str(id), (ltx, lty - 10), cv.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
    #             cv.rectangle(image, (ltx, lty), (rdx, rdy), (255,0,0), 3)
    #         output_image = os.path.join(output_test_path + "/image", image_file)
    #         print("output_image: ", output_image)
    #         cv.imwrite(output_image, image)



